#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import random

sys.path.append('.')

"""
   对FaceID合并后进行重排
"""


def main():
    parser = argparse.ArgumentParser(description="FaceID Merge ReSort")
    parser.add_argument(
        "--faceid_file_list", default="/rootfs/media/yery/Kaso/data/VideoGenderCropPerson/file_faceid_merge_list.txt", help="path to file list", type=str
    )
    parser.add_argument(
        "--faceid_merge_file_list", default="/rootfs/media/yery/Kaso/data/VideoGenderCropPerson/SelectFace/RFaceIDBoxFeatureMerge.norm.txt", help="path to faceid file list", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/yery/Kaso/data/VideoGenderCropPerson/FaceID", help="path to output", type=str
    )
    args = parser.parse_args()

    random.seed(123456)

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    faceid = 1
    merge_faceid_videos = {}
    merge_video_faceids = {}
    merge_video_files = {}
    merge_faceid_name_ids = {}
    with open(args.faceid_merge_file_list, 'r') as file:
        for line in file.readlines():
            # 5269 5269 gd/zs/sqqjd/yxxhjt/150100414a54443452067fa6d4b06600/video/20190923_152901/00023_000.jpg 1
            lines = line.strip().split()
            faceid_name = '{}_{}'.format(lines[0], lines[1])
            file_name = lines[2]
            video_name = os.path.dirname(file_name)
            cur_faceid = merge_faceid_name_ids.get(faceid_name, None)
            if cur_faceid is None:
                merge_faceid_name_ids[faceid_name] = faceid
                cur_faceid = faceid
                merge_faceid_videos[cur_faceid] = {video_name}
                faceid += 1
            else:
                merge_faceid_videos[cur_faceid].add(video_name)
            merge_video_faceids[video_name] = cur_faceid
            files = merge_video_files.get(video_name, None)
            if files is None:
                merge_video_files[video_name] = [file_name]
            else:
                files.append(file_name)

    faceid_videos = {}
    video_faceids = {}
    video_files = {}
    faceid_name_ids = {}
    with open(args.faceid_file_list, 'r') as file:
        for line in file.readlines():
            # gd/zs/sqqjd/ycjy/150100414a54443452064d2743393600/video/20190924_101121/00003_000.jpg 2599
            lines = line.strip().split()
            file_name = lines[0]
            faceid_name = '{}_{}'.format(lines[1], lines[1])
            video_name = os.path.dirname(file_name)
            cur_faceid = faceid_name_ids.get(faceid_name, None)
            if cur_faceid is None:
                faceid_name_ids[faceid_name] = faceid
                cur_faceid = faceid
                faceid_videos[cur_faceid] = {video_name}
                faceid += 1
            else:
                faceid_videos[cur_faceid].add(video_name)
            video_faceids[video_name] = cur_faceid
            files = video_files.get(video_name, None)
            if files is None:
                video_files[video_name] = [file_name]
            else:
                files.append(file_name)

    print('Video Count:', len(merge_video_faceids.keys()), len(video_faceids.keys()))
    print('FaceID Count:', len(merge_faceid_videos.keys()), len(faceid_videos.keys()))

    done_video_set = set()
    for video_name, faceid in video_faceids.items():
        if video_name in done_video_set:
            continue
        video_names = faceid_videos[faceid]
        merge_faceid_set = set()
        for video_name in video_names:
            _faceid = merge_video_faceids.get(video_name, None)
            if _faceid is None:
                continue
            merge_faceid_set.add(_faceid)
        merge_faceid_count = len(merge_faceid_set)
        if merge_faceid_count == 0:
            for video_name in video_names:
                merge_video_faceids[video_name] = faceid
                videos = merge_faceid_videos.get(faceid, None)
                if videos is None:
                    merge_faceid_videos[faceid] = {video_name}
                else:
                    videos.add(video_name)
                merge_video_files[video_name] = video_files[video_name]
        elif merge_faceid_count == 1:
            _faceid = list(merge_faceid_set)[0]
            for video_name in video_names:
                if video_name in merge_video_faceids.keys():
                    continue
                merge_video_faceids[video_name] = _faceid
                merge_faceid_videos[_faceid].add(video_name)
                merge_video_files[video_name] = video_files[video_name]
        else:
            print(video_name)

    merge_file_name_faceid = []
    merge_typical_file_name_faceid = []
    merge_one_typical_file_name_faceid = []
    typical_file_name_faceid = []
    faceid = 1
    faceid_count = 0
    video_count = 0
    for _faceid, video_names in merge_faceid_videos.items():
        video_names = list(video_names)
        video_names.sort()
        merge_files = []
        for video_name in video_names:
            video_count += 1
            merge_file_names = merge_video_files[video_name]
            merge_files.extend(merge_file_names)
            file_name = random.choice(merge_file_names)
            merge_typical_file_name_faceid.append([file_name, faceid])
            for file_name in merge_file_names:
                merge_file_name_faceid.append([file_name, faceid])

            file_names = video_files[video_name]
            file_name = random.choice(file_names)
            typical_file_name_faceid.append([file_name, faceid])
        file_name = random.choice(merge_files)
        merge_one_typical_file_name_faceid.append([file_name, faceid])
        faceid += 1
        faceid_count += 1

    merge_file_name_faceid.sort(key=lambda x: x[0])
    out_file_path = os.path.join(output_dir, 'merge_file_faceid_list.txt')
    with open(out_file_path, 'w') as file:
        for file_name, faceid in merge_file_name_faceid:
            out_info = '{} {}\n'.format(file_name, faceid)
            file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_path))

    merge_one_typical_file_name_faceid.sort(key=lambda x: x[0])
    out_file_path = os.path.join(output_dir, 'merge_one_typical_file_faceid_list.txt')
    with open(out_file_path, 'w') as file:
        for file_name, faceid in merge_one_typical_file_name_faceid:
            out_info = '{} {}\n'.format(file_name, faceid)
            file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_path))

    merge_typical_file_name_faceid.sort(key=lambda x: x[0])
    out_file_path = os.path.join(output_dir, 'merge_typical_file_faceid_list.txt')
    with open(out_file_path, 'w') as file:
        for file_name, faceid in merge_typical_file_name_faceid:
            out_info = '{} {}\n'.format(file_name, faceid)
            file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_path))

    typical_file_name_faceid.sort(key=lambda x: x[0])
    out_file_path = os.path.join(output_dir, 'typical_file_faceid_list.txt')
    with open(out_file_path, 'w') as file:
        for file_name, faceid in typical_file_name_faceid:
            out_info = '{} {}\n'.format(file_name, faceid)
            file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_path))

    print('FaceID Count:', faceid_count, ', Video Count:', video_count)
    print('finish!')


if __name__ == '__main__':
    main()
